import os
import torchvision.transforms as transforms
import torchvision.datasets as datasets
from torch.utils.data import DataLoader
from torchvision.datasets import ImageNet
from .utils import RandomSizedRectCrop, RectScale, RandomErasing, Random2DTranslation
class Data:
    def __init__(self, opt):
        pin_memory = True

        scale_size = 224

        traindir = './database/Pandora18k/train'
        valdir = './database/Pandora18k/val'
        normalize = transforms.Normalize((0.5071, 0.4865, 0.4409), (0.2673, 0.2564, 0.2762))

        trainset = datasets.ImageFolder(
            traindir,
            transforms.Compose([
                transforms.Resize(scale_size),
                RandomSizedRectCrop(224, 224),
                # transforms.RandomResizedCrop(224, scale=(0.08, 1.0)),
                transforms.RandomHorizontalFlip(),
                transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),  # 颜色抖动
                transforms.RandomRotation(15),  # 添加随机旋转
                RandomErasing(),  # 随机擦除
                Random2DTranslation(224, 224, p=0.5),  # 随机二维平移
                transforms.ToTensor(),
                normalize,
                # RectScale(224, 224),
            ]))

        self.trainLoader = DataLoader(
            trainset,
            batch_size=opt.train_batch_size,
            shuffle=True,
            num_workers=0,
            pin_memory=pin_memory,
            drop_last=True, )

        testset = datasets.ImageFolder(
            valdir,
            transforms.Compose([
                transforms.Resize(scale_size),
                transforms.CenterCrop(224),
                transforms.ToTensor(),
                normalize,
            ]))

        self.testLoader = DataLoader(
            testset,
            batch_size=opt.test_batch_size,
            shuffle=False,
            num_workers=0,
            pin_memory=pin_memory,
            drop_last=False, )

